#Script for quickly loading data for coding and demonstration purposes.

import os
import sys
import readchn
import pickle
from matplotlib import pyplot as plt

def load_data():
    #A = readchn.esc_run(run_path = 'data/example/esc/',verbose=True)
    A = readchn.esc_run(run_path = 'data/140228E2.ZIP',verbose=True)
    A.initial_find_peaks()
    A.smart_guess_peaks(0.12)
    #A.smart_fit_peaks()
    A.fit_peaks_lmfit()
    print "Done with smart_fit"
    A.compute_alpha_counts()
    return A
    

if __name__ == '__main__':
    try:
        print "Is there an existing file: ./data/quickie.pkl?"
        A = pickle.load(open('data/quickie.pkl','r'))
        print " Yes!, data object A is ready"
        print " If you want to rerun find_peaks() ... compute_alpha_counts() delete data/quickie.pkl"
    except:
        print "Could not find an existing pickle"
        print " Attempting to create a quickie.pkl"
        A = load_data()
        #A.fit_single('Po218')
        pickle.dump(A,open('data/quickie.pkl','w'))
        print " Pickled data for next time, bonne dater"



# I'm removing things that do not belong in the esc_run class to here

#Makes a heat plot of the spectra of the type specified """
"""
fig = plt.figure(figsize=(15,8), dpi=100)
ax1 = fig.add_subplot(111)
cax1 = plt.imshow(self.data, interpolation='none', aspect='auto')
pylab.colorbar()
plt.savefig('%sheat_plot_sum.png'%disc)
fig = plt.figure(figsize=(15,8), dpi=100)
ax1 = fig.add_subplot(111)
cax2 = plt.imshow(A.data, interpolation='none', aspect='auto', norm=LogNorm(vmin=1,vmax=A.data.max()))
pylab.colorbar()
plt.savefig('%sheat_plot_sum_log.png'%disc)

fig = plt.figure(figsize=(15,8), dpi=100)
ax = fig.add_subplot(111)
cax3 = plt.imshow(A.idata, interpolation='none', aspect='auto')
pylab.colorbar()
plt.savefig('%sheat_plot_ind.png'%disc)
fig = plt.figure(figsize=(15,8), dpi=100)
ax = fig.add_subplot(111)
cax4 = plt.imshow(A.idata, interpolation='none', aspect='auto', norm=LogNorm(vmin=1,vmax=self.idata.max()))
pylab.colorbar()
plt.savefig('%sheat_plot_ind_log.png'%disc)
"""


"""
   def set_energy_scale(self,energy_parms):
        self.energy_parms = energy_parms
        # This is just for makeing a plot
        if self.verbose:
            chns = numpy.arange(0,self.sum_spec.size)
            plt.subplot(2,1,1)
            plt.plot(self.sum_spec)
            plt.xlim(0,self.sum_spec.size)
            plt.xlabel('Channel #')
            plt.ylabel('Counts')

            plt.title('Linear Energy Fit y=%0.3fx+%0.3f'%(self.energy_parms[0],self.energy_parms[1]))
            plt.subplot(2,1,2)
            for i,alpha in enumerate(self.alpha_list):
                x = decays[alpha][0]['energy']
                y = max(self.sum_spec)
                plt.text(x,y,alpha,rotation='vertical')
                plt.vlines(x,0,max(self.sum_spec),colors='g',linestyles='dashed',label=alpha)
            x = (chns - self.energy_parms[1])/self.energy_parms[0]
            plt.plot(x,self.sum_spec)
            plt.xlim(4,10)
            plt.title('Smart Guess Peaks')
            plt.xlabel('Energy (MeV)')
            plt.ylabel('Counts')
            plt.show()
"""
